package com.lst.selfaiagent.advisors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 自定义 Advisor
 * CallAdvisor：用于处理同步请求和响应（非流式）
 * StreamAdvisor：用于处理流式请求和响应
 */

public class MyCustomAdvisor implements CallAdvisor, StreamAdvisor {

    private static final Logger logger = LoggerFactory.getLogger(SimpleLoggerAdvisor.class);

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain chain) {

        logger.info("开始非流式处理。。。。。");

        logger.debug("BEFORE: {}", chatClientRequest);

        ChatClientResponse chatClientResponse = chain.nextCall(chatClientRequest);

        logger.info("AFTER: {}", chatClientResponse);

        return chatClientResponse;
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain chain) {

        logger.info("开始流式处理。。。。。");
        logger.debug("BEFORE: {}", chatClientRequest);

        // 前置处理
        List<Message> messages = new ArrayList<>(chatClientRequest.prompt().getInstructions());
        messages.add(new Message() {
            @Override
            public String getText() {
                return "你需要说已经执行了前置处理";
            }

            @Override
            public Map<String, Object> getMetadata() {
                return Map.of();
            }

            @Override
            public MessageType getMessageType() {
                return MessageType.ASSISTANT;
            }
        });

        ChatClientRequest modifiedRequest = ChatClientRequest.builder()
                .prompt(new Prompt(messages, chatClientRequest.prompt().getOptions()))
                .context(chatClientRequest.context())
                .build();

        Flux<ChatClientResponse> chatClientResponses = chain.nextStream(modifiedRequest);

        return new MessageAggregator().aggregate(chatClientResponses.map(ChatClientResponse::chatResponse),
                chatResponse -> logger.debug("AFTER: {}", chatResponse))
                .map(chatResponse -> ChatClientResponse.builder().chatResponse(chatResponse).build());
    }
    
    @Override
    public String getName() {
        return "自定义的Advisor";
    }

    @Override
    public int getOrder() {
        // 默认优先级
        return 0;
    }
}